from langchain.agents import initialize_agent, AgentType
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.agents import Tool
from Tools.retrieval_qa_chatglm.bilud_chatglm3_qa_langchain import qa_chain, ChatGLMWrapper

# 将问答链封装成工具tool
def build_course_qa_tool() -> Tool:

    return Tool(
        name="CourseKnowledgeBase",
        func=qa_chain.run,
        description="用于回答课程《TensorFlow.js应用开发》相关问题。"
    )

def build_agent():
    # 加载模型
    llm = ChatGLMWrapper(
        model_path=r"E:\projects\PycharmProjects\cnsoft\Intelligent Education Agent\Tools\retrieval_qa_chatglm\models\chatglm3-6b")

    # 自定义 Prompt 模板
    custom_prompt = PromptTemplate.from_template(
        """你是一个智能教育助手，帮助用户理解课程内容。尽可能详细和准确地回答问题。

当前对话历史:
{chat_history}

用户问题:
{input}
"""
    )

    # 记忆
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

    # 构建 Agent
    agent_executor = initialize_agent(
        agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
        tools=[build_course_qa_tool()],
        llm=llm,
        memory=memory,
        verbose=True,
        agent_kwargs={"prefix": custom_prompt.template}
    )

    return agent_executor
